import numpy as np
import pandas as pd
import os
from dateutil.relativedelta import relativedelta
# import warnings
# warnings.filterwarnings('ignore')

def ratio_split(data, ratio):
    split_index = int(data.shape[0] * ratio)
    data_in = data[:split_index]
    data_out = data[split_index:]
    return data_in, data_out

def make_root(root):
    if not os.path.exists(root):
        os.makedirs(root)


class Rolling_eval_split_month():
    def __init__(self, feature, label, times, args) -> None:  
        self.args = args
        self.times = times
        self.feature = feature
        self.label = label
        self.time_slide_num = 0
        self.out_sample_side = self.args.time_param['outsample_beg'] + relativedelta(months=self.time_slide_num)
        self.out_month_list = []
        
    @property
    def out_month_num(self):
        length_m = self.args.time_param['outsample_end'].month - self.args.time_param['insample_end'].month
        length_y = self.args.time_param['outsample_end'].year - self.args.time_param['insample_end'].year
        length = length_y * 12 + length_m -1
        return length
      
    def get_beg_end(self):
        self.out_sample_side = self.args.time_param['outsample_beg'] + relativedelta(months=self.time_slide_num)
        print(f'\n当前样本外月份为:{self.out_sample_side}')
        out_beg_id = np.where(self.times>=pd.to_datetime(self.out_sample_side))[0][0]
        out_end_id = np.where(self.times>=pd.to_datetime(self.out_sample_side+ relativedelta(months=self.args.rolling_step)))[0]
        out_end_id = -1 if len(out_end_id)<=0 else out_end_id[0]    
        in_beg_id = np.where(self.times>=pd.to_datetime(self.args.time_param['insample_beg']+ relativedelta(months=self.time_slide_num)))[0][0]        
        self.out_month_list.append(self.out_sample_side)
        self.time_slide_num += self.args.rolling_step
        self.out_month_root = f'{self.args.result_root}/{str(self.out_sample_side)[:-9]}'
        if not os.path.exists(self.out_month_root):
            os.makedirs(self.out_month_root)
        return out_beg_id, out_end_id, in_beg_id
            
    def cum_split(self): 
        out_beg_id, out_end_id, _ = self.get_beg_end()     
        feature_in = self.feature[:out_beg_id-self.args.label_range]
        feature_out = self.feature[out_beg_id:out_end_id]
        label_in = self.label[:out_beg_id-self.args.label_range]
        label_out = self.label[out_beg_id:out_end_id]
        times_in = self.times[:out_beg_id-self.args.label_range]
        times_out = self.times[out_beg_id:out_end_id]
        feature_train, feature_eval = ratio_split(feature_in, self.args.train_ratio)
        label_train, label_eval = ratio_split(label_in, self.args.train_ratio)
        times_train, times_eval = ratio_split(times_in, self.args.train_ratio)
        print(f'当前样本内数据长度：{ len(times_in)},当前样本外数据长度：{ len(times_out)}')
        return feature_train, feature_eval, feature_out, label_train, label_eval, label_out, times_train, times_eval, times_out
    
    def rolling_split(self):
        out_beg_id, out_end_id, in_beg_id = self.get_beg_end()     
        feature_in = self.feature[in_beg_id:out_beg_id-self.args.label_range]
        feature_out = self.feature[out_beg_id:out_end_id]
        label_in = self.label[in_beg_id:out_beg_id-self.args.label_range]
        label_out = self.label[out_beg_id:out_end_id]
        times_in = self.times[in_beg_id:out_beg_id-self.args.label_range]
        times_out = self.times[out_beg_id:out_end_id]
        feature_train, feature_eval = ratio_split(feature_in, self.args.train_ratio)
        label_train, label_eval = ratio_split(label_in, self.args.train_ratio)
        times_train, times_eval = ratio_split(times_in, self.args.train_ratio)
        print(f'当前样本内数据长度：{ len(times_in)}, 当前样本外数据长度：{ len(times_out)}')
        return feature_train, feature_eval, feature_out, label_train, label_eval, label_out, times_train, times_eval, times_out
    

